# 数据流追踪器 - 详细记录Agent间的数据传递过程
import json
import logging
from datetime import datetime
from typing import Dict, List, Any
import os

class DataFlowTracker:
    """专门用于追踪Agent间数据传递的工具类"""

    def __init__(self, log_file="data_flow.log"):
        self.log_file = log_file
        self.flow_records = []

        # 配置专门的数据流日志
        self.logger = logging.getLogger('DataFlowTracker')
        self.logger.setLevel(logging.INFO)

        # 创建文件处理器
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(logging.INFO)

        # 创建格式化器
        formatter = logging.Formatter(
            '%(asctime)s | %(levelname)s | %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S'
        )
        file_handler.setFormatter(formatter)

        # 添加处理器到logger
        self.logger.addHandler(file_handler)

        # 同时输出到控制台
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        console_formatter = logging.Formatter(
            '%(asctime)s | 📊 DATA FLOW | %(message)s',
            datefmt='%H:%M:%S'
        )
        console_handler.setFormatter(console_formatter)
        self.logger.addHandler(console_handler)

        self.logger.info("🔄 数据流追踪器已启动")

    def log_data_flow(self, from_agent: str, to_agent: str, data: Dict[str, Any],
                     flow_type: str = "collaboration", description: str = ""):
        """记录数据流"""
        record = {
            "timestamp": datetime.now().isoformat(),
            "from_agent": from_agent,
            "to_agent": to_agent,
            "flow_type": flow_type,
            "description": description,
            "data": data
        }

        self.flow_records.append(record)

        # 格式化日志输出
        data_summary = self._summarize_data(data)
        log_message = f"{from_agent} → {to_agent} | {flow_type.upper()} | {description}"
        if data_summary:
            log_message += f" | 数据: {data_summary}"

        self.logger.info(log_message)

        # 如果是协商相关，额外记录详细内容
        if flow_type == "negotiation":
            self._log_negotiation_details(record)

    def log_agent_input(self, agent_name: str, input_data: Dict[str, Any],
                       source: str = "coordinator"):
        """记录Agent接收到的输入"""
        self.log_data_flow(
            from_agent=source,
            to_agent=agent_name,
            data=input_data,
            flow_type="input",
            description=f"{agent_name}接收输入数据"
        )

    def log_agent_output(self, agent_name: str, output_data: Dict[str, Any],
                        destination: str = "coordinator"):
        """记录Agent发送的输出"""
        self.log_data_flow(
            from_agent=agent_name,
            to_agent=destination,
            data=output_data,
            flow_type="output",
            description=f"{agent_name}返回处理结果"
        )

    def log_negotiation_request(self, from_agent: str, to_agent: str,
                               negotiation_data: Dict[str, Any]):
        """记录协商请求"""
        self.log_data_flow(
            from_agent=from_agent,
            to_agent=to_agent,
            data=negotiation_data,
            flow_type="negotiation_request",
            description=f"{from_agent}向{to_agent}发起协商"
        )

    def log_negotiation_response(self, from_agent: str, to_agent: str,
                                response_data: Dict[str, Any]):
        """记录协商响应"""
        self.log_data_flow(
            from_agent=from_agent,
            to_agent=to_agent,
            data=response_data,
            flow_type="negotiation_response",
            description=f"{from_agent}响应{to_agent}的协商请求"
        )

    def log_data_transformation(self, agent_name: str, original_data: Dict[str, Any],
                               transformed_data: Dict[str, Any], transformation_type: str):
        """记录数据转换过程"""
        transformation_record = {
            "timestamp": datetime.now().isoformat(),
            "agent": agent_name,
            "transformation_type": transformation_type,
            "original_data": original_data,
            "transformed_data": transformed_data
        }

        self.logger.info(f"🔄 {agent_name} | 数据转换 | {transformation_type}")
        self.logger.info(f"   原始数据: {self._summarize_data(original_data)}")
        self.logger.info(f"   转换后: {self._summarize_data(transformed_data)}")

    def _summarize_data(self, data: Dict[str, Any]) -> str:
        """生成数据摘要"""
        if not isinstance(data, dict):
            return str(data)[:100]

        summary_parts = []

        # 提取关键字段
        if "type" in data:
            summary_parts.append(f"type:{data['type']}")

        if "agent" in data:
            summary_parts.append(f"agent:{data['agent']}")

        if "user_requirements" in data:
            req = data["user_requirements"]
            if isinstance(req, dict):
                location = req.get("location", req.get("city", "unknown"))
                summary_parts.append(f"location:{location}")

        if "budget_data" in data:
            bd = data["budget_data"]
            if isinstance(bd, dict):
                budget = bd.get("base_budget", "unknown")
                status = bd.get("budget_status", "unknown")
                summary_parts.append(f"budget:{budget},status:{status}")

        if "weather_data" in data:
            wd = data["weather_data"]
            if isinstance(wd, dict):
                city = wd.get("city", "unknown")
                current = wd.get("current", "")[:30]
                summary_parts.append(f"weather:{city},{current}")

        if "recommendations" in data:
            recs = data["recommendations"]
            if isinstance(recs, list):
                summary_parts.append(f"recommendations:{len(recs)}项")

        if "negotiation_points" in data:
            points = data["negotiation_points"]
            if isinstance(points, list):
                summary_parts.append(f"negotiation_points:{len(points)}项")

        return " | ".join(summary_parts) if summary_parts else "复杂数据结构"

    def _log_negotiation_details(self, record: Dict[str, Any]):
        """详细记录协商过程"""
        data = record["data"]

        if "negotiation_context" in data:
            context = data["negotiation_context"]
            self.logger.info(f"   协商议题: {context.get('issue', 'unknown')}")
            self.logger.info(f"   协商提案: {context.get('proposal', 'unknown')}")

        if "proposals" in data:
            proposals = data["proposals"]
            if isinstance(proposals, list) and proposals:
                self.logger.info(f"   提案数量: {len(proposals)}")
                for i, prop in enumerate(proposals[:3]):  # 只显示前3个
                    self.logger.info(f"   提案{i+1}: {prop.get('topic', 'unknown')}")

    def get_flow_summary(self) -> Dict[str, Any]:
        """获取数据流摘要"""
        if not self.flow_records:
            return {"total_flows": 0, "agents_involved": [], "flow_types": {}}

        agents = set()
        flow_types = {}

        for record in self.flow_records:
            agents.add(record["from_agent"])
            agents.add(record["to_agent"])

            flow_type = record["flow_type"]
            flow_types[flow_type] = flow_types.get(flow_type, 0) + 1

        return {
            "total_flows": len(self.flow_records),
            "agents_involved": list(agents),
            "flow_types": flow_types,
            "time_span": {
                "start": self.flow_records[0]["timestamp"] if self.flow_records else None,
                "end": self.flow_records[-1]["timestamp"] if self.flow_records else None
            }
        }

    def print_flow_sequence(self):
        """打印数据流序列"""
        if not self.flow_records:
            print("📊 无数据流记录")
            return

        print("\n" + "="*80)
        print("🔄 AGENT间数据传递序列")
        print("="*80)

        for i, record in enumerate(self.flow_records, 1):
            timestamp = datetime.fromisoformat(record["timestamp"]).strftime("%H:%M:%S")
            from_agent = record["from_agent"]
            to_agent = record["to_agent"]
            flow_type = record["flow_type"]
            description = record["description"]

            print(f"{i:2d}. [{timestamp}] {from_agent:15} → {to_agent:15} | {flow_type:15} | {description}")

            # 显示数据摘要
            data_summary = self._summarize_data(record["data"])
            if data_summary and len(data_summary) > 50:
                print(f"      数据摘要: {data_summary[:50]}...")
            elif data_summary:
                print(f"      数据摘要: {data_summary}")

        print("="*80)
        print(f"📊 总计: {len(self.flow_records)} 次数据传递")
        print("="*80 + "\n")

# 全局数据流追踪器实例
global_tracker = DataFlowTracker()

def get_global_tracker() -> DataFlowTracker:
    """获取全局数据流追踪器"""
    return global_tracker